rm(list = ls())
set.seed(12345)
library(bpCausal)
library(readstata13)
library(ggplot2)


## change wd
## setwd("/Users/lichengl/Desktop/BSA_replication_code/")

source("code/plot_function.R")
source("code/summary_function.R")



data <- read.dta13("data/annual_data.dta")
Yname <- "topitaxrate2"
Dname <- "himobpopyear2pl1"
Xname <- c("topitaxrate2l1", "unisuffragel1", "leftexec2l1", "rgdppcl1")
index <- c("ccode", "year")

data <- data[, c(index, Yname, Dname, Xname)]
data <- na.omit(data)


# panelView(topitaxrate2 ~ himobpopyear2pl1 + unisuffragel1 + leftexec2l1 + rgdppcl1, data = data, 
#           index = c("ccode", "year"))



## naive case: w/o bsa
fit1 <-  bpCausal(data = data, index = index, 
                  Yname = Yname, Dname = Dname, Xname = Xname, 
                  Zname = NULL, Aname = NULL, re = "both", 
                  ar1 = FALSE, r = 10, niter = 25000, burn = 15000, 
                  xlasso = 1, zlasso = 1, alasso = 1, flasso = 1, 
                  a1 = 0.001, a2 = 0.001, b1 = 0.001, b2 = 0.001, 
                  c1 = 0.001, c2 = 0.001, p1 = 0.001, p2 = 0.001, 
                  sense = 0,  c12 = 0.25, c22 = 0.25, 
                  c32 = 0.5, c42 = 0.25, c52 = 0.25)




## bsa
fit12 <-  bpCausal(data = data, index = index, 
	              Yname = Yname, Dname = Dname, Xname = Xname, 
	              Zname = NULL, Aname = NULL, re = "both", 
                  ar1 = FALSE, r = 10, niter = 25000, burn = 15000, 
                  xlasso = 1, zlasso = 1, alasso = 1, flasso = 1, 
                  a1 = 0.001, a2 = 0.001, b1 = 0.001, b2 = 0.001, 
                  c1 = 0.001, c2 = 0.001, p1 = 1.001, p2 = 0.001, 
                  sense = 1, c12 = 0.25, c22 = 0.25, 
                  c32 = 0.5, c42 = 0.25, c52 = 0.25)



length(unique(c(fit12$sigma2)))/10000
## [1] 0.8165
length(unique(c(fit12$betau)))/10000
## [1] 0.6699






## condition on betau or ld and then re-run BSA

## deciles of beta_u
betau.seq <- quantile(fit12$betau, probs = seq(from = 0.1, to = 0.9, by = 0.1))
## declies of lambda_d
ld.seq <- quantile(fit12$ld, probs = seq(from = 0.1, to = 0.9, by = 0.1))


## save the distribution of lambda_d given beta_u
sim.ld <- c()

## save att conditional on a particular value of beta_u
betau_catt <- matrix(NA, 9, 4)

for (i in 1:9) {

    fit13 <-  bpCausal(data = data, index = index, 
                  Yname = Yname, Dname = Dname, Xname = Xname, 
                  Zname = NULL, Aname = NULL, re = "both", 
                  ar1 = FALSE, r = 10, niter = 25000, burn = 15000, 
                  xlasso = 1, zlasso = 1, alasso = 1, flasso = 1, 
                  a1 = 0.001, a2 = 0.001, b1 = 0.001, b2 = 0.001, 
                  c1 = 0.001, c2 = 0.001, p1 = 1.001, p2 = 0.001, 
                  sense = 1, c12 = 0.25, c22 = 0.25, 
                  c32 = 0.5, c42 = 0.25, c52 = 0.25, betau0 = betau.seq[i])


    ## 
    avg1 <- effSummary(fit13, datt = FALSE) 
    betau_catt[i, ] <- c(betau.seq[i], mean(avg1$est.avg), quantile(avg1$est.avg, c(0.025, 0.975)))

    sim.ld <- c(sim.ld, list(fit13$ld))

}

colnames(betau_catt) <- c("betau", "att", "ci_l", "ci_u")




## save the distribution of beta_u given lambda_d
sim.betau <- c()

## save att conditional on a particular value of lambda_d
ld_catt <- matrix(NA, 9, 4)


for (i in 1:9) {

    fit14 <-  bpCausal(data = data, index = index, 
                  Yname = Yname, Dname = Dname, Xname = Xname, 
                  Zname = NULL, Aname = NULL, re = "both", 
                  ar1 = FALSE, r = 10, niter = 25000, burn = 15000, 
                  xlasso = 1, zlasso = 1, alasso = 1, flasso = 1, 
                  a1 = 0.001, a2 = 0.001, b1 = 0.001, b2 = 0.001, 
                  c1 = 0.001, c2 = 0.001, p1 = 1.001, p2 = 0.001, 
                  sense = 1, c12 = 0.25, c22 = 0.25, 
                  c32 = 0.5, c42 = 0.25, c52 = 0.25, ld0 = ld.seq[i])


    avg1 <- effSummary(fit14, datt = FALSE) 
    ld_catt[i, ] <- c(ld.seq[i], mean(avg1$est.avg), quantile(avg1$est.avg, c(0.025, 0.975)))

    sim.betau <- c(sim.betau, list(fit14$betau))

}

colnames(ld_catt) <- c("ld", "att", "ci_l", "ci_u")



#save(data, fit12, fit1, betau_catt, ld_catt, sim.ld, sim.betau,
#    file = "results/SS_2012.RData")



#load("results/SS_2012.RData")


## -------------------------------------------------------------------------- ##
## 1. plot of conditional vs marginal distributions of sensitivity parameters ##
## -------------------------------------------------------------------------- ##


## 1.1 posterior dependency: ld vs betau

#v.sim.ld <- unlist(sim.ld)
## combined with original 
#v.sim.ld <- c(v.sim.ld, rep(Inf, 10000), c(fit12$ld))  ## note: we add a gap between the deciles and marginal
#index.ld <- rep(1:11, each = 10000)

#pdf("plot/exp_ld_dist.pdf", width = 10, height = 7)

#boxplot(v.sim.ld ~ index.ld, ylim = c(-3, 2), main = "",
#        ylab = expression(lambda[d]), xlab = expression(beta[u]), xaxt='n') 
#axis(1, at = c(1:9, 11), labels = c(round(betau_catt[, 1], digits = 3), ""), las = 1)

#abline(h = quantile(c(fit12$ld), 0.25), col = "red", lty = 2)
#abline(h = quantile(c(fit12$ld), 0.75), col = "red", lty = 2)

#dev.off()



## 1.2 posterior dependency: betau vs ld

#pdf("plot/exp_betau_dist.pdf", width = 10, height = 7)

#v.sim.betau <- unlist(sim.betau)
## combined with original 
#v.sim.betau <- c(v.sim.betau, rep(Inf, 10000), c(fit12$betau)) ## note: we add a gap between the deciles and marginal
#index.betau <- rep(1:11, each = 10000)

#boxplot(v.sim.betau ~ index.betau, ylim = c(-4, 4), main = "", 
#        xlab = expression(lambda[d]), ylab = expression(beta[u]), xaxt = 'n')
#axis(1, at = c(1:9, 11), labels = c(round(ld_catt[, 1], digits = 3), ""), las = 1)

#abline(h = quantile(c(fit12$betau), 0.25), col = "red", lty = 2)
#abline(h = quantile(c(fit12$betau), 0.75), col = "red", lty = 2)

#dev.off()





## -------------------------------------------------------------------------- ##
##                       2. plot of treatment effects                         ##
## -------------------------------------------------------------------------- ##


## 2.0 data preparation
## naive: 
avg1 <- effSummary(fit1)

## BSA
avg2 <- effSummary(fit12)

betau <- fit12$betau
ld <- fit12$ld 
att0 <- mean(avg1$est.avg)





## 2.1 conditional treatment effects by deciles of beta_u

y.adj <- 3
ylim <- c(0, 8)
y.seq <- ylim[1]:ylim[2]

ff <- hist(betau, breaks = 100)
ff$counts <- ff$counts / max(ff$counts)



## plot

## pdf("plot/exp_att_betau.pdf", width = 8, height = 7)

pdf("plot/figure3_left.pdf", width = 8, height = 7)
plot(ff, col = "grey", ylim = ylim, xaxt = 'n', yaxt = "n", 
         xlab = expression(beta[u]), ylab = "ATT", main = "",
         cex.lab = 1.5, cex.axis = 1.5, cex.main = 1.5, cex.sub = 1.5) 

## axis creation 
axis(1, at = betau_catt[, 1], labels = round(betau_catt[, 1], digits = 3), las = 1)
axis(2, at = y.seq, labels = y.seq - y.adj, las = 2, cex.axis = 1.5) 

points(betau_catt[, 1], betau_catt[, 2] + y.adj, pch = 19)

arrows(betau_catt[, 1], 
       betau_catt[, 3] + y.adj, 
       betau_catt[, 1], 
       betau_catt[, 4] + y.adj, 
       length = 0.1, angle = 90, code = 3)

abline(h = 0 + y.adj , lty = 3, lwd = 2, col = "red")
abline(h = att0 + y.adj, lty = "dashed", col = "blue", lwd = 2)

dev.off()







## 2.2 conditional treatment effects by deciles of lambda_d


y.adj <- 3
ylim <- c(0, 8)
y.seq = ylim[1]:ylim[2]

ff <- hist(ld, breaks = 100)
ff$counts <- ff$counts / max(ff$counts)

barplot(ff$counts,
        horiz = TRUE,                 # <- rotate bars
        names.arg = round(ff$mids, 2),  # label y-axis with bin centers
        las = 1,                      # make axis labels horizontal
        xlab = "Count",
        ylab = "Value",
        main = "Rotated Histogram")


## plot

## pdf("plot/exp_att_lambda_d.pdf", width = 8, height = 7)
pdf("plot/figure3_right.pdf", width = 8, height = 7)
plot(ff, col = "grey", ylim = ylim, xaxt = 'n', yaxt = "n", 
         xlab = expression(lambda[d]), ylab = "ATT", main = "",
         cex.lab = 1.5, cex.axis = 1.5, cex.main = 1.5, cex.sub = 1.5)

## axis creation 
axis(1, at = ld_catt[, 1], labels = round(ld_catt[, 1], digits = 3), las = 1)
axis(2, at = y.seq, labels = y.seq - y.adj, las = 2, cex.axis = 1.5) 

points(ld_catt[, 1], ld_catt[, 2] + y.adj, pch = 19)

arrows(ld_catt[, 1], 
       ld_catt[, 3] + y.adj, 
       ld_catt[, 1], 
       ld_catt[, 4] + y.adj, 
    length = 0.1, angle = 90, code = 3)

abline(h = 0 + y.adj , lty = 3, lwd = 2, col = "red")
abline(h = att0 + y.adj, lty = "dashed", col = "blue", lwd = 2)

dev.off()



## 2.2.1 contour plot for zero-effect 

att_naive <- mean(avg1$est.avg)

## samples of beta_u
bu_limit <- quantile(fit12$betau, probs = c(0.025, 0.975))
bu.seq <- seq(bu_limit[1], bu_limit[2], length.out = 100)
## samples of lambda_d
lda_limit <- quantile(fit12$ld, probs = c(0.025, 0.975))
lda.seq <- seq(lda_limit[1], lda_limit[2], length.out = 100)

att_adj <- outer(bu.seq, lda.seq, function(x, y) {att_naive - x * y})


## pdf("plot/exp_att_contour.pdf", width = 8, height = 7)
pdf("plot/figure4.pdf", width = 8, height = 7)
contour(bu.seq, lda.seq, att_adj, nlevels = 30, 
        xlab = expression(beta[u]), ylab = expression(lambda[d]),
        cex.axis = 1.5, cex.lab = 1.5, col = "black", lwd = 1.5, labcex = 1.5,
        main = "")

# Add z = 0 level in red
contour(bu.seq, lda.seq, att_adj, levels = 0, add = TRUE,
        labcex = 1.5, col = "red", lwd = 2)

att_output <- round(att_naive, 3)
# Add text annotation
points(0, 0, pch = 19, cex = 1.2, col = "black")
text(0, 0, labels = paste("Unadjusted \n (", att_output, ")", sep = ""), pos = 3, cex = 1.2, col = "black")

dev.off()



## 2.3 compare naive, bsa, different level of lambda, ld 

## pdf("plot/exp_att_compare.pdf", width = 8, height = 7)
pdf("plot/figure2_right.pdf", width = 8, height = 7)
plot(c(0, 1, 2, 3), 
     c(NA, mean(avg1$est.avg), mean(avg2$est.avg), NA),
     ylim = c(-2, 5),
     pch = 19, xlab = "", ylab = "ATT",
     main = "", xaxt = "n", 
     cex.lab = 1.5, cex.axis = 1.5, cex.main = 1.5, cex.sub = 1.5)

# hack: we draw arrows but with very special "arrowheads"
arrows(c(0, 1, 2, 3), 
        c(NA, quantile(avg1$est.avg, 0.025), quantile(avg2$est.avg, 0.025), NA), 
        c(0, 1, 2, 3), 
        c(NA, quantile(avg1$est.avg, 0.975), quantile(avg2$est.avg, 0.975), NA), 
        length = 0.1, angle = 90, code = 3)

abline(h = 0, lty = 3, lwd = 2, col = "red")

axis(1, at = c(1, 2), 
     labels = c("w/o BSA", "w/ BSA"), 
     col.axis = "black", cex.axis = 1.5, cex.lab = 1.5)

dev.off()





## 2.4 dynamic ATT w/ and w/o BSA

TT <- dim(avg2$est.eff)[1]
pos <- which(rownames(avg2$est.eff) == "0")
period <- as.numeric(rownames(avg2$est.eff))

pos <- which(period >= -10 & period <= 6)
period.plot <- period[pos]


## r plot 
## pdf("plot/exp_datt.pdf", width = 8, height = 7)
pdf("plot/figure2_left.pdf", width = 8, height = 7)
plot(period.plot, avg2$est.eff[pos, 5], col = "red", 
     type = "l", lwd = 5, ylim = c(-6, 15), 
     xlab = "Year relative to Treatment", ylab = "ATT", main = "", 
     cex.lab=1.5, cex.axis=1.5, cex.main=1.5, cex.sub=1.5)

lines(period.plot, avg2$est.eff[pos, 6], col = "red", lwd = 3, lty = 3)
lines(period.plot, avg2$est.eff[pos, 7], col = "red", lwd = 3, lty = 3)

lines(period.plot, avg1$est.eff[pos, 5], col = "blue", lwd = 5)
lines(period.plot, avg1$est.eff[pos, 6], col = "blue", lwd = 3, lty = 3)
lines(period.plot, avg1$est.eff[pos, 7], col = "blue", lwd = 3, lty = 3)

abline(v = 0, lty = 3, lwd = 2)
abline(h = 0, lty = 3, lwd = 2)

legend(-10, 15, legend = c("W/ BSA", "W/o BSA"),
       col = c("red", "blue"), lty = 1:1, lwd = 2:2, pt.cex = 1.5)

dev.off()





























